# -*- coding: utf-8 -*-

import tushare as ts
import pymysql
from sqlalchemy import create_engine
from urllib.parse import quote_plus
from datetime import datetime
import pandas as pd
import matplotlib.pyplot as plt
import mplfinance as mpf

# ... existing code ...
# 配置 MySQL 连接信息
# %40 表示 @ 符号，%23 表示 # 符号
password = quote_plus('NewPass123!@#')
engine = create_engine(f'mysql+pymysql://root:{password}@10.54.1.6:13306/zjjk')

# 设置 tushare 的 token
pro = ts.pro_api('b1125846b73e4e0445d2d8cd1ea3d7ba1443e881c78eb09acb78f6ec')

def plot_stock_close_price():
    """
    从 tushare_002929 表取出近 60 行记录，以 trade_date 倒序排序，
    然后绘制 trade_date 为横坐标、close 为纵坐标的曲线。
    """
    # 从数据库获取数据
    table_name = 'tushare_002929_sz'
    query = f"SELECT trade_date, open, high, low, close, vol FROM {table_name} ORDER BY trade_date DESC LIMIT 120" 
    data = pd.read_sql(query, engine)

    # 将 trade_date 转换为日期类型并升序排序
    data['trade_date'] = pd.to_datetime(data['trade_date'])
    data = data.sort_values('trade_date')
    data = data.set_index('trade_date')
    data.rename(columns={'vol': 'volume'}, inplace=True)

    # 设置图片清晰度
    plt.rcParams['figure.dpi'] = 300

    # 计算布林带数据
    mpf_kwargs = dict(
        type='candle',
        style='yahoo',
        volume=True,
        title='002929 近120日K线图',
        ylabel='价格',
        ylabel_lower='成交量',
        figratio=(12, 6),
        addplot=[
            mpf.make_addplot(data['close'].rolling(window=20).mean(), color='b', label='MA20'),
            mpf.make_addplot(data['close'].rolling(window=20).mean() + 2 * data['close'].rolling(window=20).std(), color='g', label='Upper Band'),
            mpf.make_addplot(data['close'].rolling(window=20).mean() - 2 * data['close'].rolling(window=20).std(), color='r', label='Lower Band')
        ]
    )

    # 绘制 K 线图和布林轨道
    mpf.plot(data, **mpf_kwargs)

if __name__ == '__main__':
    end_date = datetime.now().strftime('%Y%m%d')
    plot_stock_close_price()  # 调用绘图函数